import "dotenv/config";
import readline from "readline";
import { ChatOpenAI } from "@langchain/openai";
import { ChatMessageHistory } from "langchain/memory";
import { getBufferString } from "@langchain/core/messages";
import { ChatPromptTemplate } from "@langchain/core/prompts";
import {
  RunnableSequence,
  RunnablePassthrough,
} from "@langchain/core/runnables";
import { StringOutputParser } from "@langchain/core/output_parsers";

// 模型
const chatModel = new ChatOpenAI({
  apiKey: process.env.OPENAI_API_KEY,
  configuration: {
    baseURL: process.env.OPENAI_API_BASE,
  },
});

// prompt
const chatPrompt = ChatPromptTemplate.fromMessages([
  [
    "system",
    `You are a helpful assistant. Answer all questions to the best of your ability.

Here is the chat history summary:
{history_summary}`,
  ],
  ["human", "{input}"],
]);

// summary chain
const summaryPrompt = ChatPromptTemplate.fromMessages([
  ["system", "请总结以下对话内容：\n{new_lines}\n\n已有摘要：\n{summary}"],
]);

const summaryChain = RunnableSequence.from([
  summaryPrompt,
  chatModel,
  new StringOutputParser(),
]);

// memory
const history = new ChatMessageHistory();
let summary = "";

// 主 chatChain
const chatChain = RunnableSequence.from([
  new RunnablePassthrough({
    func: async (input) => {
      await history.addUserMessage(input);
      return input;
    },
  }),
  RunnablePassthrough.assign({
    history_summary: () => summary,
  }),
  chatPrompt,
  chatModel,
  new StringOutputParser(),
  new RunnablePassthrough({
    func: async (output) => {
      await history.addAIChatMessage(output);
      const messages = await history.getMessages();
      const new_lines = getBufferString(messages);
      summary = await summaryChain.invoke({
        summary,
        new_lines,
      });
      return output;
    },
  }),
]);

// CLI loop
const rl = readline.createInterface({
  input: process.stdin,
  output: process.stdout,
});

console.log("💬 开始对话吧！（输入 exit 可退出）");

const loop = () => {
  rl.question("> 你：", async (input) => {
    if (input.trim().toLowerCase() === "exit") {
      rl.close();
      return;
    }

    try {
      const res = await chatChain.invoke({ input });
      console.log("🤖 AI：", res);
      console.log("🧠 当前摘要：", summary);
    } catch (err) {
      console.error("❌ 错误：", err);
    }

    loop();
  });
};

loop();
